import torch.nn as nn
from torch import cat

"""
GoogleNet
"""

class GoogleNet(nn.Module):
    def __init__(self, classes):
        super().__init__()
        self.__create_nn(classes)

    def __create_nn(self, classes):
        #-------------------------------------------------------------------------------#
        # 第一块主干网络
        #-------------------------------------------------------------------------------#
        self.block1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=3),
            nn.ReLU(True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1)),
            nn.ReLU(True),
            nn.Conv2d(64, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
            nn.ReLU(True),
            nn.MaxPool2d(3, 2)
        )
        # -------------------------------------------------------------------------------#
        # Inception3a
        # -------------------------------------------------------------------------------#
        self.inception3a = Inception(192, 64, 96, 128, 16, 32, 32)
        # -------------------------------------------------------------------------------#
        # Inception2
        # -------------------------------------------------------------------------------#
        self.inception3b = Inception(256, 128, 128, 192, 32, 96, 64)
        # -------------------------------------------------------------------------------#
        # 第二块主干网络-池化
        # -------------------------------------------------------------------------------#
        self.block2 = nn.Sequential(
            nn.MaxPool2d(3, 2, padding=1)
        )
        # -------------------------------------------------------------------------------#
        # Inception4a
        # -------------------------------------------------------------------------------#
        self.inception4a = Inception(480, 192, 96, 208, 16, 48, 64)
        # -------------------------------------------------------------------------------#
        # Aux Net1
        # -------------------------------------------------------------------------------#
        self.aux1 = InceptionAux([512, 128, 2048, 1024, classes])
        # -------------------------------------------------------------------------------#
        # Inception4b
        # -------------------------------------------------------------------------------#
        self.inception4b = Inception(512, 160, 112, 224, 24, 64, 64)
        # -------------------------------------------------------------------------------#
        # Inception4c
        # -------------------------------------------------------------------------------#
        self.inception4c = Inception(512, 128, 128, 256, 24, 64, 64)
        # -------------------------------------------------------------------------------#
        # Inception4d
        # -------------------------------------------------------------------------------#
        self.inception4d = Inception(512, 112, 144, 288, 32, 64, 64)
        # -------------------------------------------------------------------------------#
        # Inception4e
        # -------------------------------------------------------------------------------#
        self.inception4e = Inception(528, 256, 160, 320, 32, 128, 128)
        # -------------------------------------------------------------------------------#
        # Aux Net2
        # -------------------------------------------------------------------------------#
        self.aux2 = InceptionAux([528, 128, 2048, 1024, classes])
        # 第三块主干网络-池化
        # -------------------------------------------------------------------------------#
        self.block3 = nn.Sequential(
            nn.MaxPool2d(3, 2, padding=1)
        )
        # -------------------------------------------------------------------------------#
        # Inception5a
        # -------------------------------------------------------------------------------#
        self.inception5a = Inception(832, 256, 160, 320, 32, 128, 128)
        # -------------------------------------------------------------------------------#
        # Inception5b
        # -------------------------------------------------------------------------------#
        self.inception5b = Inception(832, 384, 192, 384, 48, 128, 128)
        # 第四块主干网络-池化
        # -------------------------------------------------------------------------------#
        self.block4 = nn.Sequential(
            nn.AvgPool2d(7, 1),
            nn.Flatten(),
            nn.Dropout(0.4),
            nn.Linear(1024, classes),
            nn.Softmax(dim=1)
        )
        self.__init_weight()

    def __init_weight(self):
        for module in self.modules():
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
            elif isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, 0, 0.01)
                nn.init.constant_(module.bias, 0)

    def forward(self, x):
        x = self.block1(x)
        # 3
        x = self.inception3a(x)
        x = self.inception3b(x)
        # block2
        x = self.block2(x)
        # 4a
        x = self.inception4a(x)
        # aux1
        if self.training:
            aux1 = self.aux1(x)
        # 4b
        x = self.inception4b(x)
        # 4c
        x = self.inception4c(x)
        # 4d
        x = self.inception4d(x)
        # aux2
        if self.training:
            aux2 = self.aux2(x)
        # 4e
        x = self.inception4e(x)
        # block3
        x = self.block3(x)
        # 5
        x = self.inception5a(x)
        x = self.inception5b(x)
        # block4
        x = self.block4(x)

        if self.training:
            return x, aux1, aux2
        return x

"""
Inception
"""

class Inception(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.__inception(*args)

    def __inception(self, input_channel, channel_121, channel_323_reduce, channel_323, channel_525_reduce, channel_525, channel_pool_121):
        self.block_one = nn.Sequential(
            # 512 -> 128
            nn.Conv2d(input_channel, channel_121, kernel_size=(1, 1), stride=(1, 1)),
            nn.ReLU(True)
        )
        self.block_two = nn.Sequential(
            # 512 -> 256
            nn.Conv2d(input_channel, channel_323_reduce, kernel_size=(1, 1), stride=(1, 1)),
            nn.ReLU(True),
            # 256 -> 128
            nn.Conv2d(channel_323_reduce, channel_323, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(True)
        )
        self.block_three = nn.Sequential(
            # 512 -> 24
            nn.Conv2d(input_channel, channel_525_reduce, kernel_size=(1, 1), stride=(1, 1)),
            nn.ReLU(True),
            # 24 -> 64
            nn.Conv2d(channel_525_reduce, channel_525, kernel_size=(5, 5), stride=(1, 1), padding=2),
            nn.ReLU(True)
        )
        self.block_four = nn.Sequential(
            nn.MaxPool2d(3, 1, padding=1),
            # 512 -> 64
            nn.Conv2d(input_channel, channel_pool_121, kernel_size=(1, 1), stride=(1, 1)),
            nn.ReLU(True)
        )

    def forward(self, x):
        block1 = self.block_one(x)
        block2 = self.block_two(x)
        block3 = self.block_three(x)
        block4 = self.block_four(x)
        x = cat([block1, block2, block3, block4], dim=1)
        return x

"""
辅助分类器
"""
class InceptionAux(nn.Module):
    def __init__(self, arg):
        super().__init__()
        self.__create_nn(*arg)

    def __create_nn(self, in_channel, out_channel, in_feature, out_feature, classes):
        self.sequential = nn.Sequential(
            nn.AvgPool2d(5, 3, padding=1),
            nn.Conv2d(in_channel, out_channel, kernel_size=(1, 1), stride=(1, 1)),
            nn.ReLU(True),
            nn.Flatten(),
            nn.Dropout(0.5),
            nn.Linear(in_feature, out_feature),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(out_feature, classes),
            nn.Softmax(dim=1)
        )


    def forward(self, x):
        return self.sequential(x)

if __name__ == "__main__":
    from utils.utils import get_model_flops_args
    get_model_flops_args(GoogleNet(5), (3, 224, 224))